"""
增强版测试数据生成脚本
运行此脚本生成测试数据：python test_data.py
"""
import sqlite3
import datetime
import random
import os

def create_enhanced_test_data():
    """创建增强版测试数据"""
    # 如果数据库存在，先备份
    if os.path.exists('accounting.db'):
        backup_name = f"accounting_backup_{datetime.datetime.now().strftime('%Y%m%d_%H%M%S')}.db"
        os.rename('accounting.db', backup_name)
        print(f"已备份原数据库: {backup_name}")
    
    conn = sqlite3.connect('accounting.db')
    cursor = conn.cursor()
    
    # 初始化数据库表
    cursor.execute('''
        CREATE TABLE IF NOT EXISTS transactions (
            id INTEGER PRIMARY KEY AUTOINCREMENT,
            type TEXT NOT NULL,
            amount REAL NOT NULL,
            main_category TEXT NOT NULL,
            sub_category TEXT NOT NULL,
            item_detail TEXT,
            note TEXT,
            date TEXT NOT NULL,
            time TEXT NOT NULL,
            created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
        )
    ''')
    
    cursor.execute('''
        CREATE TABLE IF NOT EXISTS budgets (
            id INTEGER PRIMARY KEY AUTOINCREMENT,
            category TEXT NOT NULL,
            amount REAL NOT NULL,
            period TEXT NOT NULL,
            created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
        )
    ''')
    
    cursor.execute('''
        CREATE TABLE IF NOT EXISTS budget_alerts (
            id INTEGER PRIMARY KEY AUTOINCREMENT,
            threshold INTEGER NOT NULL DEFAULT 80,
            is_enabled BOOLEAN NOT NULL DEFAULT 0,
            created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
        )
    ''')
    
    # 插入默认预算预警设置
    cursor.execute('SELECT COUNT(*) FROM budget_alerts')
    if cursor.fetchone()[0] == 0:
        cursor.execute('''
            INSERT INTO budget_alerts (threshold, is_enabled)
            VALUES (80, 0)
        ''')
    
    # 清空现有数据
    cursor.execute('DELETE FROM transactions')
    cursor.execute('DELETE FROM budgets')
    
    # 设置预算数据
    budgets = [
        ('食品酒水', 2000, 'monthly'),
        ('交通出行', 800, 'monthly'),
        ('衣服饰品', 1000, 'monthly'),
        ('居家物业', 2500, 'monthly'),
        ('休闲娱乐', 600, 'monthly'),
        ('学习进修', 500, 'monthly'),
        ('医疗保健', 300, 'monthly'),
        ('人情往来', 400, 'monthly'),
        ('其他支出', 200, 'monthly')
    ]
    
    for budget in budgets:
        cursor.execute('''
            INSERT INTO budgets (category, amount, period)
            VALUES (?, ?, ?)
        ''', budget)
    
    # 定义详细的分类和项目
    expense_categories = {
        "食品酒水": {
            "餐饮": ["早餐", "午餐", "晚餐", "外卖", "快餐", "餐厅"],
            "食材": ["蔬菜", "水果", "肉类", "水产", "蛋奶", "粮油", "调味品"],
            "零食饮料": ["零食", "饮料", "咖啡", "茶饮", "甜品", "冰淇淋"],
            "烟酒": ["香烟", "白酒", "啤酒", "红酒", "其他酒类"]
        },
        "交通出行": {
            "公共交通": ["公交", "地铁", "出租车", "网约车", "共享单车"],
            "私家车": ["加油", "停车费", "洗车", "保养", "保险", "年检"],
            "长途交通": ["火车", "飞机", "长途汽车", "船票"]
        },
        "衣服饰品": {
            "服装": ["上衣", "裤子", "裙子", "外套", "内衣", "睡衣"],
            "鞋帽": ["鞋子", "帽子", "袜子", "围巾"],
            "饰品": ["首饰", "手表", "眼镜", "皮带", "包包"]
        },
        "居家物业": {
            "房租房贷": ["房租", "房贷"],
            "水电物业": ["水费", "电费", "燃气费", "物业费", "取暖费"],
            "家居用品": ["家具", "家电", "厨具", "清洁用品", "收纳用品"]
        },
        "休闲娱乐": {
            "娱乐活动": ["电影", "KTV", "游乐场", "演出", "展览"],
            "旅游度假": ["住宿", "景点门票", "旅游团费"],
            "运动健身": ["健身房", "瑜伽", "游泳", "球类运动"]
        },
        "学习进修": {
            "书籍资料": ["图书", "杂志", "电子书", "学习资料"],
            "教育培训": ["学费", "培训费", "考证费", "辅导班"]
        },
        "医疗保健": {
            "看病买药": ["门诊费", "药费", "检查费", "手术费"],
            "保健养生": ["保健品", "体检", "理疗", "按摩"]
        },
        "人情往来": {
            "礼物": ["生日礼物", "节日礼物", "结婚礼物", "纪念日礼物"],
            "红包": ["现金红包", "微信红包", "支付宝红包"],
            "请客": ["请朋友", "请同事", "请家人", "商务请客"]
        },
        "其他支出": {
            "其他": ["意外支出", "未分类支出", "手续费"]
        }
    }
    
    income_categories = {
        "职业收入": {
            "工资": ["基本工资", "绩效工资", "年终奖", "津贴"],
            "奖金": ["项目奖金", "销售提成", "季度奖金"],
            "兼职": ["兼职收入", "自由职业", "咨询服务"]
        },
        "投资理财": {
            "投资收益": ["股票", "基金", "理财产品", "外汇"],
            "利息收入": ["存款利息", "债券利息", "P2P利息"]
        },
        "其他收入": {
            "退款": ["购物退款", "服务退款", "违约金"],
            "红包": ["微信红包", "支付宝红包", "现金红包"],
            "其他": ["中奖", "二手出售", "报销"]
        }
    }
    
    # 生成从2024年到现在的数据
    end_date = datetime.datetime.now()
    start_date = datetime.datetime(2024, 1, 1)
    
    current_date = start_date
    transaction_id = 0
    
    print("正在生成测试数据...")
    
    while current_date <= end_date:
        date_str = current_date.strftime("%Y-%m-%d")
        
        # 工作日通常有更多交易
        is_weekday = current_date.weekday() < 5
        transaction_count = random.randint(2, 8) if is_weekday else random.randint(1, 4)
        
        # 每月的第一天可能有房租等固定支出
        is_first_day = current_date.day == 1
        
        # 生成支出记录
        for _ in range(transaction_count):
            # 选择随机的支出分类
            main_cat = random.choice(list(expense_categories.keys()))
            middle_cat = random.choice(list(expense_categories[main_cat].keys()))
            sub_cat = random.choice(expense_categories[main_cat][middle_cat])
            
            # 生成金额（根据分类有不同的金额范围）
            amount_ranges = {
                "食品酒水": (5, 100),
                "交通出行": (2, 50),
                "衣服饰品": (50, 500),
                "居家物业": (100, 3000),
                "休闲娱乐": (20, 300),
                "学习进修": (30, 1000),
                "医疗保健": (10, 500),
                "人情往来": (50, 1000),
                "其他支出": (5, 200)
            }
            
            min_amount, max_amount = amount_ranges.get(main_cat, (5, 100))
            
            # 如果是每月第一天，可能会有房租等大额支出
            if is_first_day and main_cat == "居家物业" and middle_cat == "房租房贷":
                amount = random.uniform(1500, 3000)
            else:
                amount = round(random.uniform(min_amount, max_amount), 2)
            
            time_str = f"{random.randint(8, 22):02d}:{random.randint(0, 59):02d}:{random.randint(0, 59):02d}"
            
            # 生成具体项目和备注
            item_details = {
                "早餐": ["包子", "豆浆", "油条", "面包", "牛奶"],
                "午餐": ["米饭", "面条", "炒菜", "快餐", "自助餐"],
                "晚餐": ["火锅", "烧烤", "炒菜", "西餐", "日料"],
                "蔬菜": ["青菜", "西红柿", "土豆", "胡萝卜", "黄瓜"],
                "水果": ["苹果", "香蕉", "橙子", "葡萄", "西瓜"],
                "公交": ["上班", "下班", "外出", "回家"],
                "地铁": ["上班", "下班", "外出", "回家"],
                "加油": ["92号", "95号", "98号", "柴油"]
            }
            
            item_detail = random.choice(item_details.get(sub_cat, [""]))
            notes = ["日常消费", "必需品", "娱乐", "工作相关", "家庭用品", "个人护理"]
            note = random.choice(notes)
            
            cursor.execute('''
                INSERT INTO transactions 
                (type, amount, main_category, sub_category, item_detail, note, date, time)
                VALUES (?, ?, ?, ?, ?, ?, ?, ?)
            ''', ('expense', amount, main_cat, f"{middle_cat} - {sub_cat}", 
                  item_detail, note, date_str, time_str))
        
        # 每月10号和25号可能有工资收入
        if current_date.day in [10, 25]:
            main_cat = "职业收入"
            middle_cat = "工资"
            sub_cat = "基本工资"
            amount = round(random.uniform(5000, 15000), 2)
            time_str = "09:00:00"
            note = "工资收入"
            
            cursor.execute('''
                INSERT INTO transactions 
                (type, amount, main_category, sub_category, item_detail, note, date, time)
                VALUES (?, ?, ?, ?, ?, ?, ?, ?)
            ''', ('income', amount, main_cat, f"{middle_cat} - {sub_cat}", 
                  "", note, date_str, time_str))
        
        # 偶尔有其他收入
        if random.random() < 0.05:  # 5%的概率
            main_cat = random.choice(list(income_categories.keys()))
            middle_cat = random.choice(list(income_categories[main_cat].keys()))
            sub_cat = random.choice(income_categories[main_cat][middle_cat])
            
            amount_ranges = {
                "职业收入": (100, 5000),
                "投资理财": (50, 2000),
                "其他收入": (10, 500)
            }
            
            min_amount, max_amount = amount_ranges.get(main_cat, (10, 500))
            amount = round(random.uniform(min_amount, max_amount), 2)
            time_str = f"{random.randint(9, 18):02d}:{random.randint(0, 59):02d}:{random.randint(0, 59):02d}"
            
            cursor.execute('''
                INSERT INTO transactions 
                (type, amount, main_category, sub_category, item_detail, note, date, time)
                VALUES (?, ?, ?, ?, ?, ?, ?, ?)
            ''', ('income', amount, main_cat, f"{middle_cat} - {sub_cat}", 
                  "", "其他收入", date_str, time_str))
        
        # 每月处理一次提交，避免内存占用过多
        if current_date.day % 10 == 0:
            conn.commit()
        
        current_date += datetime.timedelta(days=1)
    
    conn.commit()
    conn.close()
    
    # 统计生成的数据
    conn = sqlite3.connect('accounting.db')
    cursor = conn.cursor()
    
    cursor.execute('SELECT COUNT(*) FROM transactions')
    total_count = cursor.fetchone()[0]
    
    cursor.execute('SELECT COUNT(*) FROM transactions WHERE type = "income"')
    income_count = cursor.fetchone()[0]
    
    cursor.execute('SELECT COUNT(*) FROM transactions WHERE type = "expense"')
    expense_count = cursor.fetchone()[0]
    
    cursor.execute('SELECT MIN(date), MAX(date) FROM transactions')
    min_date, max_date = cursor.fetchone()
    
    conn.close()
    
    print("测试数据生成完成！")
    print(f"数据统计:")
    print(f"- 总交易记录: {total_count} 条")
    print(f"- 收入记录: {income_count} 条")
    print(f"- 支出记录: {expense_count} 条")
    print(f"- 数据时间范围: {min_date} 至 {max_date}")
    print(f"- 预算设置: {len(budgets)} 个分类")

def analyze_test_data():
    """分析测试数据的分布情况"""
    conn = sqlite3.connect('accounting.db')
    cursor = conn.cursor()
    
    print("\n数据分布分析:")
    
    # 按年份统计
    cursor.execute('SELECT strftime("%Y", date) as year, COUNT(*) FROM transactions GROUP BY year ORDER BY year')
    year_stats = cursor.fetchall()
    print("按年份分布:")
    for year, count in year_stats:
        print(f"  {year}年: {count} 条记录")
    
    # 按月份统计（最近一年）
    cursor.execute('''
        SELECT strftime("%Y-%m", date) as month, COUNT(*) 
        FROM transactions 
        WHERE date >= date('now', '-1 year')
        GROUP BY month 
        ORDER BY month
    ''')
    month_stats = cursor.fetchall()
    print("最近12个月分布:")
    for month, count in month_stats:
        print(f"  {month}: {count} 条记录")
    
    # 按分类统计支出
    cursor.execute('''
        SELECT main_category, COUNT(*), SUM(amount) 
        FROM transactions 
        WHERE type = 'expense'
        GROUP BY main_category 
        ORDER BY SUM(amount) DESC
    ''')
    category_stats = cursor.fetchall()
    print("支出分类统计:")
    for category, count, total in category_stats:
        print(f"  {category}: {count} 条记录, 总金额 {total:.2f} 元")
    
    # 按类型统计
    cursor.execute('SELECT type, COUNT(*), SUM(amount) FROM transactions GROUP BY type')
    type_stats = cursor.fetchall()
    print("收支类型统计:")
    for t, count, total in type_stats:
        type_name = "收入" if t == "income" else "支出"
        print(f"  {type_name}: {count} 条记录, 总金额 {total:.2f} 元")
    
    conn.close()

if __name__ == "__main__":
    create_enhanced_test_data()
    analyze_test_data()